'''
Created on Mar 22, 2012

@author: pedro
'''

import Preprocessor as pre
import HDRGenerators
import HDRNormalizers
import Metaprocessing
import sys
import os
import pickle
import numpy as np
import time

from Network import OscillatorPyramidalLayer,OutputLayer
from Neurons import PyramidalNeuron,OutputNeuron,integrate_timecourse
from Utils import angle

from random import shuffle,randint
from multiprocessing import Process,current_process



def hdr_battery_heatmap(battery,filename):
    battery_arr = range(len(battery))    

    for key, value in battery.iteritems():
        battery_arr[key] = value
    
    pts = [range(len(battery_arr)) for x in range(len(battery_arr))]
 
    count = len(battery_arr) * (len(battery_arr)-1)
    total = 0.0
 
    for x in range(len(battery_arr)):
        for y in range(len(battery_arr)):
            diff = angle(battery_arr[x],battery_arr[y]) if x != y else 0.0
            total += diff
            pts[x][y] = diff
    
    for x in range(len(battery_arr)):
        for y in range(len(battery_arr)):
            if pts[x][y] == 0:
                pts[x][y] = (total)/(count)
    
    if filename != None:
        Metaprocessing.save_to_heatmap(pts,filename)
        
    return pts


def init_hdr_battery(hdr_file="", config_params={}):
    '''Initializes an HDR battery, ready for injecting into neural network
    
    Keyword Args:
    hdr_file      --- filename of HDR. No filename results in a generated HDR.
    config_params --- Configuration parameters dictionary
    
    @return - an HDR battery either generated from params or loaded from a file
    '''
    battery = None
    usigmoid_intercept      = config_params["usigmoid_intercept"]
    usigmoid_slope          = config_params["usigmoid_slope"]
    usigmoid_gaussian_mu    = config_params["usigmoid_gaussian_mu"]
    usigmoid_gaussian_sigma = config_params["usigmoid_gaussian_sigma"]
    
    try:
        battery = pre.load_battery_from_file(hdr_file)
        print "Loaded battery from file, found {0} HDRs".format(len(battery))
    except:
        print "Could not find battery file, creating a new one"
        gen = HDRGenerators.UniformSigmoidStarHDRGenerator(usigmoid_intercept,
                                                           usigmoid_slope,
                                                           usigmoid_gaussian_mu,
                                                           usigmoid_gaussian_sigma)
        norm = HDRNormalizers.FloorHDRNormalizer()
        battery = pre.generate_hdr_battery(config_params["num_hdrs"]+1,
                                           config_params["numPyrs"],
                                           gen,
                                           norm)
        pre.save_battery_to_file(battery,"pre_fired_battery.dat")
        
    return battery


def init_neural_net(num_outputs, config_params, pickled_network=None):
    '''Initialize neural network layers using parameters.
    
    @returns - a two-tuple of the (1) Pyramidal Layer, and (2) Output Layer
    '''
    v0              = config_params["v0"]
    Rm              = config_params["Rm"]
    Cm              = config_params["Cm"]
    tau_ref         = config_params["tau_ref"]
    Vth             = config_params["Vth"]
    V_spike         = config_params["V_spike"]
    oscillator_freq = config_params["oscillator_freq"]
    max_weight      = config_params["max_weight"]
    numPyrs         = config_params["numPyrs"]
    connectivity    = config_params["connectivity"]
    stdp_decay_factor = config_params["stdp_decay_factor"]

    pyramidals = []
    outputs = []
    
    if pickled_network is None:
        for _ in xrange(num_outputs):
            o1 = OutputNeuron(v0,Rm,Cm,tau_ref,Vth,V_spike)
            outputs.append(o1)
            o1 = None
        
        for _ in xrange(numPyrs):
            pyr1 = PyramidalNeuron(v0,Rm,Cm,tau_ref,Vth,V_spike)
            pyramidals.append(pyr1)
            pyr1 = None
    
        pyrlayer = OscillatorPyramidalLayer(pyramidals, oscillator_freq)
        outlayer = OutputLayer(pyramidals, 
                               outputs, 
                               stdp_decay_factor, 
                               max_weight, 
                               connectivity)
        
        pickled_network_file = open("pickled_network.old","w+")
        pickle.dump(obj=(pyrlayer,outlayer),file=pickled_network_file)
        
    else:
        pickled_network_file = open(pickled_network,"r")
        pyramidals, outlayer = pickle.load(pickled_network_file)
                           
    return (pyrlayer, outlayer)
    

def run_simulation(config_params, hdr_file, data_file, network_pickle=None):
    
    '''Set variables from config'''
    dt              = config_params["dt"]
    sim_time        = config_params["time"]
    num_pplus_hdrs  = config_params["num_pplus_hdrs"]
    numOuts         = config_params["numPyrs"]
    v0              = config_params["v0"]
    time_stop       = config_params["time_stop"]
    num_extra_trials= config_params["num_extra_trials"]
    num_regular_trials = config_params["num_regular_trials"]
    
    '''---Initialize HDR Battery---'''  
    battery = init_hdr_battery(hdr_file, config_params)
    hdr_battery_angle_matrix = np.array(hdr_battery_heatmap(battery, None))
    np.savetxt("hdr_bat_{0}.txt".format(current_process().ident), hdr_battery_angle_matrix)
    
    
    '''---Initialize Neural Network---'''
    pyrlayer,outlayer = init_neural_net(numOuts, 
                                        config_params, 
                                        pickled_network=network_pickle)
    
    hdr_baseline_outputs = {}
    hdr_posttraining_outputs = {}
    
    '''---Choose P+ HDRs Randomly---'''
    pplus_hdrs = []
    while len(pplus_hdrs) < num_pplus_hdrs:
        attempt = randint(0,len(battery)-1)
        if attempt in pplus_hdrs:
            continue
        else:
            pplus_hdrs.append(attempt)

    
    '''---Part 1: Establish Baseline Responses---'''
    
    
    """Each HDR is presented $num_regular_trials times at random"""
    trials_sequence = []
    for i in range(len(battery)):
        for _ in range(num_regular_trials):
            trials_sequence.append(i)
    shuffle(trials_sequence)
    
    for hdr_num in trials_sequence:
        #Reset neuron's Vm to v0
        pyrlayer.reset_neurons()
        outlayer.reset_neurons()
        
        #Fire Pyramidal Layer first
        pi_layer_outputs = pyrlayer.fire(battery[hdr_num], dt, sim_time)
        
        #Fire Output Layer Last
        output_layer_outputs = outlayer.fire(pi_layer_outputs,
                                             learning=True,
                                             params_dict=config_params)
        hdr_baseline_outputs[hdr_num] = sum([integrate_timecourse(output[1], v0, time_stop) for output in output_layer_outputs])
    
    pyrlayer.reset_neurons()
    outlayer.reset_neurons()
    
    """Acquire Baseline Responses by Firing w/ Learning OFF"""
    for hdr_num in xrange(len(battery)):
        #Fire Pyramidal Layer first
        pi_layer_outputs = pyrlayer.fire(battery[hdr_num], dt, sim_time)
        
        #Fire Output Layer Last
        output_layer_outputs = outlayer.fire(pi_layer_outputs,
                                             learning=False,
                                             params_dict=config_params)
        
        hdr_baseline_outputs[hdr_num] = sum([integrate_timecourse(output[1], v0, time_stop) for output in output_layer_outputs])
        
        pyrlayer.reset_neurons()
        outlayer.reset_neurons()


    '''---Part 2: Perform Additional Training for P+ HDRs---'''
        
    for _ in xrange(num_extra_trials):
        for hdr_num in pplus_hdrs:
            #Reset neuron's Vm to v0
            pyrlayer.reset_neurons()
            outlayer.reset_neurons()
            
            #Fire Pyramidal Layer first
            pi_layer_outputs = pyrlayer.fire(battery[hdr_num], dt, sim_time)
            
            #Fire Output Layer Last
            output_layer_outputs = outlayer.fire(pi_layer_outputs,
                                                 learning=True,
                                                 params_dict=config_params)
                                                 
            hdr_posttraining_outputs[hdr_num] = sum([integrate_timecourse(output[1], v0, time_stop) for output in output_layer_outputs])
        
    pyrlayer.reset_neurons()
    outlayer.reset_neurons()
    
    
    
    '''---Part 3: Acquire Post-Training Responses by Firing w/ Learning OFF---'''
    
    for hdr_num in range(len(battery)):
        #Fire Pyramidal Layer first
        pi_layer_outputs = pyrlayer.fire(battery[hdr_num], dt, sim_time)
        
        #Fire Output Layer Last
        output_layer_outputs = outlayer.fire(pi_layer_outputs,
                                             learning=False,
                                             params_dict=config_params)
        
        hdr_posttraining_outputs[hdr_num] = sum([integrate_timecourse(output[1], v0, time_stop) for output in output_layer_outputs])
        
        pyrlayer.reset_neurons()
        outlayer.reset_neurons()
        
    bminus = []
    bplus = []
    pminus = []
    pplus = []
    output_data = []
    
    for hdr_num in range(len(battery)):
        if hdr_num in pplus_hdrs:
            output_data.append((1,hdr_num,hdr_baseline_outputs[hdr_num],hdr_posttraining_outputs[hdr_num],hdr_posttraining_outputs[hdr_num]-hdr_baseline_outputs[hdr_num]))
            bplus.append((hdr_num,hdr_baseline_outputs[hdr_num]))
            pplus.append((hdr_num,hdr_posttraining_outputs[hdr_num]))
        else:
            output_data.append((0,hdr_num,hdr_baseline_outputs[hdr_num],hdr_posttraining_outputs[hdr_num],hdr_posttraining_outputs[hdr_num]-hdr_baseline_outputs[hdr_num]))
            bminus.append((hdr_num,hdr_baseline_outputs[hdr_num]))
            pminus.append((hdr_num,hdr_posttraining_outputs[hdr_num]))
        
    print "\nBaseline v.s. Post-Training for B-/P- HDRs:"    
    for index in range(len(bminus)):
        print "\tHDR #{0} B-: {1}".format(bminus[index][0],bminus[index][1])
        print "\tHDR #{0} P-: {1}".format(pminus[index][0],pminus[index][1])
        print "\tHDR #{0} Delta: {1}".format(pminus[index][0],pminus[index][1]-bminus[index][1])
        
    print "\nBaseline v.s. Post-Training for B+/P+ HDRs:"    
    for index in range(len(bminus)):
        print "\tHDR #{0} B+: {1}".format(bplus[index][0],bplus[index][1])
        print "\tHDR #{0} P+: {1}".format(pplus[index][0],pplus[index][1])
        print "\tHDR #{0} Delta: {1}".format(pplus[index][0],pplus[index][1]-bplus[index][1])
    
    print "Done"
    
    Metaprocessing.write_to_output_file(output_data, data_file)
    
    #Pickle the current network state
    network_pickle_file = open("pickled_network.curr","w+")
    pickle.dump(obj=(pyrlayer,outlayer),file=network_pickle_file)

"""--------------------------------------------------------------------------------------------------"""

if __name__ == '__main__':
    
    config_file = sys.argv[1]
    hdr_file = sys.argv[2]
    data_file = sys.argv[3]
    
    pickled_file = None
    try:
        pickled_file = sys.argv[4]
    except Exception:
        pass
    
    '''---Load Config---'''
    config_params = Metaprocessing.get_config_params(config_file)
    
    '''---Make new directory and CD into it---'''
    old_dir = os.getcwd()
    dir_name = os.path.splitext(os.path.basename(config_file))[0]
    dir_name = dir_name + str(int(time.time()))
    os.mkdir(os.getcwd()+"/"+dir_name)
    os.chdir(os.getcwd()+"/"+dir_name)
        
    for i in range(config_params["num_iterations"]): 
        run_simulation(config_params, hdr_file, data_file, pickled_file)
            
    #Metaprocessing.get_average_hdr_heatmap()
    Metaprocessing.calculate_results_from_output(data_file,
                                                 segregate_baseline=False,
                                                 statfilename='stats.txt')
    
    #CD back into old directory
    os.chdir(old_dir)
    

    
